import asyncio
import os.path
import sys
import time
from urllib import parse

import httpx
import m3u8
from httpx import HTTPError

default_headers = {
    'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/107.0.0.0 Safari/537.36 Edg/107.0.1418.56'
}


class M3U8Downloader:
    def __init__(self, uri: str, name: str = None, sub_name: str = None, headers: dict = None):
        up = parse.urlparse(uri)
        ps = up.path.split('/')
        self.uri = uri
        self.headers = dict(**default_headers, **headers) if headers is not None else default_headers
        self.playlist = m3u8.load(uri, headers=self.headers, verify_ssl=False)
        self.name = name or ps[-2] or up.netloc or 'm3u8'
        self.sub_name = sub_name or ps[-1].split('.')[0] or time.strftime("%Y%m%d%H%M%S", time.localtime(time.time()))
        self._requests = None
        self._downloading = False
        self._sem = None
        self._total = 0
        self._finished = 0
        self._key_index = 1

    def _get_real_playlist(self):
        playlist = self.playlist
        while hasattr(playlist, 'is_variant') and playlist.is_variant:
            playlist = m3u8.load(playlist.playlists[0].absolute_uri, headers=self.headers, verify_ssl=False)
        self.playlist = playlist
        return playlist

    def _get_process(self):
        percent = 0 if self._total == 0 else round(self._finished / self._total * 100, 2)
        total_len = len(str(self._total))
        return f'[{repr(self._finished).zfill(total_len)}/{self._total} {percent:.2f}]'

    def _get_temp_path(self, path: str):
        temppath = os.path.join(path, f'{self.name}_{self.sub_name}')
        if not os.path.exists(temppath):
            os.makedirs(temppath)
        return temppath

    async def _download_key(self, key: m3u8.Key, path: str, retry=3):
        if self._requests is None:
            self._requests = httpx.AsyncClient(base_url=self.playlist.base_uri, verify=False, timeout=10)
        if retry < 1:
            retry = 1
        times = 1
        while times <= retry:
            try:
                res = await self._requests.get(key.absolute_uri, headers=self.headers)
                if res.status_code == 200:
                    filepath = self._get_temp_path(path)
                    key_file = f'key_{self._key_index}.key'
                    full_name = os.path.join(filepath, key_file)
                    with open(full_name, 'wb') as f:
                        f.write(res.content)
                    key.uri = key_file

                    print("Downloaded Key: {0}".format(key.uri))
                    self._key_index = self._key_index + 1

                    return True
            except HTTPError as e:
                print(f"{key.uri} ==> Times: {times} Error: {e}")
            except Exception as e:
                print(f"{key.uri} ==> Times: {times} Exception: {e}")
                return False
            times = times + 1
        print("{0} ==> 失败 {1} 次，跳过。".format(key.uri, times))
        return False

    async def _download_segment(self, segment: m3u8.Segment, path: str, retry=5):
        async with self._sem:
            url_path = segment.uri.split('/')[-1].split('?')[0]
            ts_type = url_path.split('.')[-1].lower()
            if ts_type in ['png', 'jpg', 'bmp', 'jpeg']:
                segment.byterange = '100000000@8'
            # 已经下载
            if segment.custom_parser_values.get('downloaded'):
                self._finished = self._finished + 1
                print("{0} Downloaded: {1}".format(self._get_process(), segment.uri))
                return True
            segment.custom_parser_values['downloaded'] = False

            # 初始化下载器
            if self._requests is None:
                self._requests = httpx.AsyncClient(base_url=self.playlist.base_uri, verify=False, timeout=10)

            if retry < 1:
                retry = 1

            times = 1
            while times <= retry:
                try:
                    res = await self._requests.get(segment.absolute_uri, headers=self.headers)
                    if res.status_code == 200:
                        filepath = self._get_temp_path(path)
                        full_name = os.path.join(filepath, url_path)
                        with open(full_name, 'wb') as f:
                            f.write(res.content)

                        segment.custom_parser_values['downloaded'] = True
                        segment.uri = url_path

                        self._finished = self._finished + 1
                        print("{0} Downloaded: {1}".format(self._get_process(), segment.uri))

                        return True
                except HTTPError as e:
                    print(f"{segment.uri} ==> Times: {times} Error: {e}")
                except Exception as e:
                    print(f"{segment.uri} ==> Times: {times} Exception: {e}")
                    return False
                times = times + 1

            print("{0} {1} ==> 失败 {2} 次，跳过。".format(self._get_process(), segment.uri, times))
            return False

    async def _download_segments(self, segments: list[m3u8.Segment], path: str):
        await asyncio.gather(*[self._download_segment(segment, path) for segment in segments])

        # 重试失败片段
        segments = list(filter(lambda x: x.custom_parser_values.get('downloaded') is False, segments))
        while len(segments) > 0:
            answer = input(f'提示：有 {len(segments)} 个片段未下载成功，是否重试？(yes or No)')
            if answer.lower() not in ('y', 'yes'):
                print('No!')
                return
            print('Yes!')
            await asyncio.gather(*[self._download_segment(segment, path) for segment in segments])
            segments = list(filter(lambda x: x.custom_parser_values.get('downloaded') is False, segments))

    def _concat_playlist(self, playlist, path: str, remove_temp=True):
        segments = list(filter(lambda x: x.custom_parser_values.get('downloaded') is False, playlist.segments))
        if len(segments) > 0:
            print(f'提示：有 {len(segments)} 个片段未下载成功，无法合成视频。')
            return

        current_path = os.getcwd()
        temppath = self._get_temp_path(path)
        os.chdir(temppath)
        filename = os.path.join('..', f'{self.name}_{self.sub_name}.mp4')
        command = f'ffmpeg -allowed_extensions ALL -protocol_whitelist "file,http,https,tls,crypto,tcp" -i "index.m3u8" -c copy "{filename}"'
        exit_code = os.system(command)
        if exit_code == 0:
            print('合并成功：' + os.path.join(path, filename))
            if remove_temp:
                os.chdir(path)
                if sys.platform == 'win32':
                    os.system(f'rmdir /s /q "{temppath}"')
                else:
                    os.system(f'rm -rf "{temppath}"')
        else:
            print(f'合并失败：{exit_code}')
        os.chdir(current_path)

    async def _download_playlist(self, path: str):
        playlist = self._get_real_playlist()

        self._total = len(playlist.segments)
        if self._total == 0:
            print('提示：没有可用的片段')
            return
        print(f'开始下载：片段数量：{self._total}')
        self._finished = 0

        for key in playlist.keys:
            if key:
                await self._download_key(key, path, 2)

        # 下载片段
        await self._download_segments(playlist.segments, path)
        # 写入本地m3u8文件
        temppath = self._get_temp_path(path)
        filepath = os.path.join(temppath, 'index.m3u8')
        playlist.dump(filepath)
        # 合成文件
        self._concat_playlist(playlist, path)

    def transfer(self, path: str, sem=16):
        if self._downloading:
            print('当前正在下载中...')
            return
        self._downloading = True

        if sem < 1:
            sem = 16
        self._sem = asyncio.Semaphore(sem)
        print("准备下载 {0} - {1}，并行任务数:{2}".format(self.name, self.sub_name, sem))
        asyncio.run(self._download_playlist(path))
        print("下载完成 {0} - {1}".format(self.name, self.sub_name))
